import requests
import json

class BaiduAISearch:
    def __init__(self, api_key):
        self.api_key = api_key
        self.base_url = "https://qianfan.baidubce.com/v2/ai_search/web_search"
    
    def search(self, query, **kwargs):
        """
        执行搜索
        
        Args:
            query: 搜索关键词
            **kwargs: 其他可选参数
        """
        headers = {
            'Content-Type': 'application/json',
            'Authorization': f'Bearer {self.api_key}'
        }
        
        # 构建请求体
        data = {
            "messages": [
                {
                    "content": query,
                    "role": "user"
                }
            ],
            "search_source": "baidu_search_v2",
            "resource_type_filter": kwargs.get('resource_type_filter', [{"type": "web", "top_k": 20}]),
            "search_recency_filter": kwargs.get('search_recency_filter', "year")
        }
        
        # 可选参数
        if 'search_filter' in kwargs:
            data['search_filter'] = kwargs['search_filter']
        
        if 'block_websites' in kwargs:
            data['block_websites'] = kwargs['block_websites']
        
        if 'edition' in kwargs:
            data['edition'] = kwargs['edition']
        
        try:
            response = requests.post(self.base_url, headers=headers, json=data, timeout=30)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            print(f"请求失败: {e}")
            return None

# 使用示例
def baidu_search(search_arg):
    
    API_KEY = "xxx"
    
    # 创建搜索实例
    searcher = BaiduAISearch(API_KEY)
    
    # 基础搜索
    result = searcher.search(search_arg)
    
    if result and 'references' in result:
        res=''
        print("搜索成功!")
        for ref in result['references']:
            res+=(f"{ref['id']}. {ref['title']}")
            res+=(f"   摘要: {ref['content'][:100]}...")

    else:
        res+=("搜索失败!")
        if result and 'code' in result:
            res+=(f"错误码: {result['code']}, 错误信息: {result['message']}")

    return res

